from sklearn.metrics import average_precision_score
import torch
import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=20):
    for epoch in range(num_epochs):
        model.train()
        train_loss = 0
        for mel_spec, labels in train_loader:
            mel_spec = mel_spec.to(device).permute(0, 2, 1)
            labels = labels.to(device)

            optimizer.zero_grad()
            outputs = model(mel_spec)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

        model.eval()
        val_loss = 0
        with torch.no_grad():
            for mel_spec, labels in val_loader:
                mel_spec = mel_spec.to(device).permute(0, 2, 1)
                labels = labels.to(device)
                outputs = model(mel_spec)
                loss = criterion(outputs, labels)
                val_loss += loss.item()

        print(
            f"Epoch [{epoch + 1}/{num_epochs}], Train Loss: {train_loss / len(train_loader):.4f}, Val Loss: {val_loss / len(val_loader):.4f}")


def evaluate_model(model, test_loader):
    model.eval()
    y_true, y_pred = [], []
    with torch.no_grad():
        for mel_spec, labels in test_loader:
            mel_spec = mel_spec.to(device).permute(0, 2, 1)
            labels = labels.cpu().numpy()
            outputs = torch.sigmoid(model(mel_spec)).cpu().numpy()

            y_true.append(labels)
            y_pred.append(outputs)

    y_true = np.vstack(y_true)
    y_pred = np.vstack(y_pred)
    map_score = average_precision_score(y_true, y_pred, average='macro')
    print(f"mAP: {map_score:.4f}")
